// Package archive 提供文件和文件夹的压缩与解压功能
// 支持格式: zip, tar, tar.gz
// 支持压缩和解压单个文件、多个文件或文件夹
package archive

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// CompressOptions 压缩选项
type CompressOptions struct {
	Format      string // 压缩格式: zip, tar, tar.gz
	Destination string // 目标文件路径
}

// Compress 压缩文件或文件夹
func Compress(sources []string, options CompressOptions) error {
	if len(sources) == 0 {
		return fmt.Errorf("no sources provided")
	}

	// 确保目标目录存在
	destDir := filepath.Dir(options.Destination)
	if err := os.MkdirAll(destDir, 0755); err != nil {
		return fmt.Errorf("failed to create destination directory: %v", err)
	}

	// 根据格式选择压缩方法
	switch strings.ToLower(options.Format) {
	case "zip":
		return compressZip(sources, options.Destination)
	case "tar":
		return compressTar(sources, options.Destination, false)
	case "tar.gz":
		return compressTar(sources, options.Destination, true)
	default:
		return fmt.Errorf("unsupported compression format: %s", options.Format)
	}
}

// compressZip 使用zip格式压缩
func compressZip(sources []string, dest string) error {
	zipFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	for _, source := range sources {
		err := filepath.Walk(source, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			// 创建文件头
			header, err := zip.FileInfoHeader(info)
			if err != nil {
				return err
			}

			// 设置压缩方法
			header.Method = zip.Deflate

			// 确保路径在压缩文件中是相对的
			relPath, err := filepath.Rel(filepath.Dir(source), path)
			if err != nil {
				return err
			}
			if info.IsDir() {
				relPath += "/"
			}
			header.Name = relPath

			// 写入文件头
			writer, err := zipWriter.CreateHeader(header)
			if err != nil {
				return err
			}

			// 如果是目录，不需要写入内容
			if info.IsDir() {
				return nil
			}

			// 打开并写入文件内容
			file, err := os.Open(path)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(writer, file)
			return err
		})

		if err != nil {
			return err
		}
	}

	return nil
}

// compressTar 使用tar格式压缩，可选gzip压缩
func compressTar(sources []string, dest string, useGzip bool) error {
	// 创建目标文件
	file, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer file.Close()

	var tw *tar.Writer
	if useGzip {
		gw := gzip.NewWriter(file)
		defer gw.Close()
		tw = tar.NewWriter(gw)
	} else {
		tw = tar.NewWriter(file)
	}
	defer tw.Close()

	for _, source := range sources {
		err := filepath.Walk(source, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			// 创建tar头
			header, err := tar.FileInfoHeader(info, info.Name())
			if err != nil {
				return err
			}

			// 确保路径在压缩文件中是相对的
			relPath, err := filepath.Rel(filepath.Dir(source), path)
			if err != nil {
				return err
			}
			header.Name = relPath

			// 写入文件头
			if err := tw.WriteHeader(header); err != nil {
				return err
			}

			// 如果是目录，不需要写入内容
			if info.IsDir() {
				return nil
			}

			// 打开并写入文件内容
			file, err := os.Open(path)
			if err != nil {
				return err
			}
			defer file.Close()

			_, err = io.Copy(tw, file)
			return err
		})

		if err != nil {
			return err
		}
	}

	return nil
}
